#
import numpy as np
import torch
import torch.nn as nn

class AttentionHead(nn.Module):
    def __init__(self, width, head_size):
        super().__init__()
        self.head_size = head_size

        self.query = nn.Linear(width, head_size)
        self.key = nn.Linear(width, head_size)
        self.value = nn.Linear(width, head_size)

    def forward(self, x, mask=None):
        # Obtaining Queries, Keys, and Values
        Q = self.query(x)
        K = self.key(x)
        V = self.value(x)

        # Dot Product of Queries and Keys
        attention = Q @ K.transpose(-2,-1)

        # Scaling
        attention = attention / (self.head_size ** 0.5)

        # Applying Attention Mask
        if mask is not None:
            attention = attention.masked_fill(mask == 0, float("-inf"))

        attention = torch.softmax(attention, dim=-1)

        attention = attention @ V

        return attention
  
class MultiHeadAttention(nn.Module):
    def __init__(self, width, n_heads):
        super().__init__()
        self.head_size = width // n_heads

        self.W_o = nn.Linear(width, width)

        self.heads = nn.ModuleList([AttentionHead(width, self.head_size) for _ in range(n_heads)])

    def forward(self, x, mask=None):
        # Combine attention heads
        out = torch.cat([head(x, mask=mask) for head in self.heads], dim=-1)

        out = self.W_o(out)

        return out